#pragma once

#include <cstdint>
#include <arpa/inet.h>

#include "network/protocols.h"

// Ethernet frame header.
struct [[gnu::packed]] EthHeader {
    std::uint8_t dstMac[6];
    std::uint8_t srcMAc[6];
    std::uint16_t ethType;

    // Type check.
    bool isIPv4() const { return ntohs(ethType) == protocols::erthertype::TPV4;}
    bool isIPv6() const { return ntohs(ethType) == protocols::erthertype::IPV6;}
    bool isVLAN() const { return ntohs(ethType) == protocols::erthertype::VLAN;}
};

struct [[gnu::packed]] IPv4Header {
    std::uint8_t versionIhl;
    std::uint8_t tos;
    std::uint16_t totalLength;
    std::uint16_t id;
    std::uint16_t flagsFragment;
    std::uint8_t ttl;
    std::uint8_t protocol;
    std::uint16_t checksum;
    std::uint32_t srcAddr;
    std::uint32_t dstAddr;

    std::uint8_t version() const { return versionIhl >> 4;}
    std::uint8_t headerLength() const { return (versionIhl & 0x0F) * 4;}

    // Type check.
    bool isTCP() const { return protocol == 6;}
    bool isUDP() const { return protocol == 17;}
    bool isICMP() const { return protocol == 1;}
};

// Simplified IPv6 header.
struct [[gnu::packed]] IPv6Header {
    std::uint32_t verTcFlow;            // version 4 bytes, traffic class 8 bytes, flow label 20 bytes
    std::uint16_t payloadLength;
    std::uint8_t nextHeader;
    std::uint8_t hopLimit;
    std::uint8_t srcAddr[16];
    std::uint8_t dstAddr[16];

    std::uint8_t version() const { return (verTcFlow >> 28) & 0x0F;}
};

struct [[gnu::packed]] TCPHeader {
    std::uint16_t srcPort;
    std::uint16_t dstPort;
    std::uint32_t seqNum;
    std::uint32_t ackNum;
    std::uint8_t dataOffset;
    std::uint8_t flags;
    std::uint16_t windowSize;
    std::uint16_t checksum;
    std::uint16_t urgentPtr;

    std::uint8_t headerLength() const { return (dataOffset >> 4) * 4;}

    // Flags check.
    bool isSYN() const { return flags & 0x02;}
    bool isACK() const { return flags & 0x10;}
    bool isFIN() const { return flags & 0x01;}
    bool isRST() const { return flags & 0x04;}
};

struct [[gnu::packed]] UDPHeader {
    std::uint16_t srcPort;
    std::uint16_t dstPort;
    std::uint16_t length;
    std::uint16_t checksum;
};